import video
from concurrent.futures import ProcessPoolExecutor


def get_indices(sr_path, gt_path, mode='RGB', crop=0, w=False):
    print(sr_path + " : " + gt_path + " start.")
    res = dict()
    res['key'] = f'{sr_path} : {gt_path}'
    res['psnr'] = video.get_psnr(sr_path, gt_path, mode=mode, crop=crop)
    res['ssim'] = video.get_ssim(sr_path, gt_path, mode=mode, crop=crop)
    if w:
        res['ws_psnr'] = video.get_ws_psnr(sr_path, gt_path, mode=mode, crop=crop)
        res['ws_ssim'] = video.get_ws_ssim(sr_path, gt_path, mode=mode, crop=crop)
    print(sr_path + " : " + gt_path + " end.")
    return res


if __name__ == '__main__':
    # 配置参数
    sr_path_list = [
        "../LbasicSR/results/BasicVSRPP_x4_MIG/visualization/REDS4_x4/5",
        "../LbasicSR/results/BasicVSRPP_x4_MIG/visualization/REDS4_x4/6",
        "../LbasicSR/results/BasicVSRPP_x4_MIG/visualization/REDS4_x4/8",
        "../LbasicSR/results/BasicVSRPP_x4_MIG/visualization/REDS4_x4/9",
    ]
    gt_path_list = [
        "../MIG/test/VR/newtest/HR/5",
        "../MIG/test/VR/newtest/HR/6",
        "../MIG/test/VR/newtest/HR/8",
        "../MIG/test/VR/newtest/HR/9",
    ]
    save_path = "./result_mig_5689.txt"
    mode = "Y"
    crop = 0
    is_w = True

    # 提交任务到进程池
    ps = []
    with ProcessPoolExecutor() as pool:
        for sr_path, gt_path in zip(sr_path_list, gt_path_list):
            ps.append(pool.submit(get_indices, sr_path, gt_path, mode, crop, is_w))
    pool.shutdown(True)
    # 保存结果
    total_psnr = 0
    total_ssim = 0
    total_ws_psnr = 0
    total_ws_ssim = 0
    for p in ps:
        res = p.result()
        with open(save_path, "a") as f:
            f.write("\n" + str(res['key']) + "\n")
            f.write("MODE:" + mode + "\n")
            f.write("CROP:" + str(crop) + "\n")
            f.write("PSNR:\t" + str(res['psnr']) + "\n")
            f.write("SSIM:\t" + str(res['ssim']) + "\n")
            total_psnr += res['psnr']
            total_ssim += res['ssim']
            if is_w:
                f.write("WS_PSNR:\t" + str(res['ws_psnr']) + "\n")
                f.write("WS_SSIM:\t" + str(res['ws_ssim']) + "\n")
                total_ws_psnr += res['ws_psnr']
                total_ws_ssim += res['ws_ssim']
            f.write("\n------------------------------------\n")

    count = len(ps)
    with open(save_path, "a") as f:
        f.write("Average Matrixs\n")
        f.write("MODE:" + mode + "\n")
        f.write("CROP:" + str(crop) + "\n")

        f.write("Average PSNR:\t" + str(total_psnr / count) + "\n")
        f.write("Average SSIM:\t" + str(total_ssim / count) + "\n")
        if is_w:
            f.write("WS_PSNR:\t" + str(total_ws_psnr / count) + "\n")
            f.write("WS_SSIM:\t" + str(total_ws_ssim / count) + "\n")
